import os
from PIL import Image
from tqdm import tqdm
import shutil

def convert_to_gray_by_class():
    # 基础路径设置
    src_dir = '/ifs/root/data/ImageNet/ImageNet2012/val'  # 原始验证集路径（包含所有JPEG文件）
    dst_base = '/ifs/root/ipa01/101/user_101002/Project/ControlVAR-main/dataset/imagenet2012/gray_images/ImageNet2012/val_gray'  # 灰度图保存的基础路径
    filelist_path = '/ifs/root/data/ImageNet/ImageNet2012/val/filelist.txt'  # filelist.txt的路径
    
    # 创建基础目标目录
    os.makedirs(dst_base, exist_ok=True)
    
    # 读取文件列表
    with open(filelist_path, 'r') as f:
        file_list = [line.strip() for line in f]
    
    # 处理每个文件
    for file_path in tqdm(file_list):
        # 解析类别目录和文件名
        class_dir, img_name = file_path.split('/')
        
        # 创建目标类别目录
        dst_class_dir = os.path.join(dst_base, class_dir)
        os.makedirs(dst_class_dir, exist_ok=True)
        
        # 设置源文件和目标文件的完整路径
        src_path = os.path.join(src_dir, img_name)  # 源文件在val目录下
        dst_path = os.path.join(dst_class_dir, img_name)
        
        # 转换并保存灰度图
        try:
            with Image.open(src_path) as img:
                gray_img = img.convert('L')
                gray_img.save(dst_path)
        except Exception as e:
            print(f"处理图片 {img_name} 时出错: {str(e)}")

if __name__ == '__main__':
    convert_to_gray_by_class()